import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import ssl

# 处理SSL证书问题
ssl._create_default_https_context = ssl._create_unverified_context

# 设备配置
if (torch.cuda.is_available()):
    device = torch.device('cuda')
elif (torch.backends.mps.is_available() and torch.backends.mps.is_built()):
    device = torch.device('mps')
else:
    device = torch.device('cpu')
print(f"使用设备: {device}")

# 超参数
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
])

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=True,
                                           transform=transform,
                                           download=True)

test_dataset = torchvision.datasets.MNIST(root='./data',
                                          train=False,
                                          transform=transform)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

# 显示样本数据
def imshow(img):
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)), cmap='gray')
    plt.show()

# 显示一些训练样本
dataiter = iter(train_loader)
images, labels = next(dataiter)
print("显示一些训练样本:")
imshow(torchvision.utils.make_grid(images[:4]))
print('标签: ', ' '.join(f'{labels[j]}' for j in range(4)))

# 定义卷积神经网络
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # 第一个卷积层: 输入通道1，输出通道16，卷积核大小5x5
        self.conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2)
        # 最大池化层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 第二个卷积层: 输入通道16，输出通道32，卷积核大小5x5
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2)
        # 全连接层
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        # 输入形状: (batch_size, 1, 28, 28)
        out = self.conv1(x)           # (batch_size, 16, 28, 28)
        out = self.relu(out)
        out = self.pool(out)          # (batch_size, 16, 14, 14)
        out = self.conv2(out)         # (batch_size, 32, 14, 14)
        out = self.relu(out)
        out = self.pool(out)          # (batch_size, 32, 7, 7)
        out = out.reshape(out.size(0), -1)  # 展平: (batch_size, 32*7*7)
        out = self.fc1(out)           # (batch_size, 128)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)           # (batch_size, 10)
        return out

# 创建模型实例
model = ConvNet().to(device)
print("\nCNN模型结构:")
print(model)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
print("\n开始训练CNN模型...")
total_step = len(train_loader)
loss_list = []
acc_list = []

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        loss_list.append(loss.item())
        
        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Loss: {loss.item():.4f}')
            
            # 计算准确率
            with torch.no_grad():
                correct = 0
                total = 0
                for images, labels in test_loader:
                    images = images.to(device)
                    labels = labels.to(device)
                    outputs = model(images)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
                
                accuracy = 100 * correct / total
                acc_list.append(accuracy)
                print(f'测试集准确率: {accuracy:.2f}%')

# 测试模型
print("\n开始测试CNN模型...")
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    
    print(f'CNN模型最终测试准确率: {100 * correct / total:.2f}%')

# 保存模型
torch.save(model.state_dict(), 'mnist_cnn_model.pth')
print("\n模型已保存为'mnist_cnn_model.pth'")

# 显示预测结果
print("\n显示CNN预测结果:")
model.eval()
with torch.no_grad():
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images_show = images
    images = images.to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)
    
    print('真实标签: ', ' '.join(f'{labels[j]}' for j in range(12)))
    print('预测标签: ', ' '.join(f'{predicted[j]}' for j in range(12)))
    
    imshow(torchvision.utils.make_grid(images_show[:12]))

# 绘制训练曲线
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(loss_list)
plt.title('CNN训练损失')
plt.xlabel('步骤')
plt.ylabel('损失')

plt.subplot(1, 2, 2)
plt.plot(acc_list)
plt.title('CNN测试准确率')
plt.xlabel('评估次数')
plt.ylabel('准确率 (%)')
plt.show()

print("\nCNN模型训练完成!")